#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "elftool.h"
#include "memsim.h"
#include "isa.h"

#define TAG_MATCH 0
#define TAG_STEP 1
#define TAG_REG 2
#define TAG_MEMREAD 3 
#define TAG_MEMWRITE 4

//////////////////////////////////////////////////////////////////
// Reading logs

FILE* logfile = 0;
char log_buffer[256];
void full_log_init(char* name) {
	char str[128];
	strcpy(str, name);
	strcat(str, ".log");

	logfile =  fopen(str, "r");
	log_buffer[0] = 0;
}

void log_read_nextline() {
	if ( !feof(logfile) ){
		fgets(log_buffer, 256, logfile);
	} else {
		log_buffer[0] = '#';
		log_buffer[1] = '\0';
	}
}
char* log_get_nextline() {
	return log_buffer;
}


//////////////////////////////////////////////////////////////////
// Execution Related

exec_section* esections;
unsigned char ecount = 0;

unsigned int fetch(unsigned int addr) {
	int i;

	for ( i = 0; i < ecount; i++ ) {
		unsigned int offset = esections[i].offset;
		unsigned int size = esections[i].size;
		if ( addr >= offset && addr < offset + size ) {
			unsigned int ret = reverse_uint(esections[i].data[(addr-offset)>>2]);
			return ret;
		}
	}

	//TODO: raise exeption?!
//	fprintf(stderr, "error: fetching invalid section(%x)\n", addr );
	return 0;
//	exit(-1);
}

signed int reg[32];

unsigned int next_pc(unsigned int pc, instr i) {
	switch ( i.opcode ) {
		case J: return i.target << 2;
		case JAL: 
			reg[31] = pc + 4;
			return i.target << 2;
		case JR:
			if ( i.fn == F_JALR) {
				reg[i.rd] = 0;//pc + 4;
				return i.rs;
			} else if ( i.fn == F_JR ) {
				return i.rs;
			}
			return 0;//pc + 4;			
		case BEQ:
			if ( reg[i.rs] == reg[i.rt] ) return pc + 4 + (i.imm<<2);
			else return 0;//pc + 4;
		case BNE:
			if ( reg[i.rs] != reg[i.rt] ) return pc + 4 + ( i.imm<<2);
			else return 0;//pc + 4;
		case BLEZ:
			if ( reg[i.rs] <= 0 ) return pc + 4 + (i.imm<<2);
			else return 0;//pc + 4;
		case BGTZ:
			if ( reg[i.rs] > 0 ) return pc + 4 + ( i.imm<<2);
			else return 0;//pc + 4;
		case BMORE:
			if (i.rd == D_BLTZ ) {
				if ( reg[i.rs] < 0 ) return pc + 4 + ( i.imm<<2);
				else return 0;//pc + 4;
			} else if (i.rd == D_BGEZ) {
				if ( reg[i.rs] >= 0 ) return pc + 4 + ( i.imm<<2);
				else return 0;//pc + 4;
			}
			return 0;
		default: return 0;
	}

	fprintf(stderr, "error: unhandled branch instruction!\n");
	return 0;
}

int alu(instr i) {
	switch(i.opcode) {
		case ADDI: reg[i.rt] = reg[i.rs] + i.imm; 
			//TODO: handle overflow!
			return 0;
		case ADDIU: reg[i.rt] = reg[i.rs] + i.imm; return 0;
		case SLTI: reg[i.rt] = reg[i.rs] < i.imm ? 1 : 0; return 0;
		case SLTIU: reg[i.rt] = reg[i.rs] < (unsigned short)i.imm ? 1 : 0; return 0;
		case ANDI: reg[i.rt] = reg[i.rs] & i.imm; return 0;
		case ORI: reg[i.rt] = reg[i.rs] | i.imm; return 0;
		case XORI: reg[i.rt] = reg[i.rs] ^ i.imm; return 0;
		case LUI: reg[i.rt] = i.imm << 16; return 0;
		case 0: {
			switch(i.fn) {
				case F_SLL: reg[i.rd] = reg[i.rt]<<i.sh; return 0;
				case F_SRL: reg[i.rd] = reg[i.rt]>>i.sh; return 0;
				case F_SRA: return 1; // TODO
				case F_SRAV: return 1; // TODO
				case F_SLLV: reg[i.rd] = reg[i.rt]<<reg[i.rs]; return 0;
				case F_SRLV: reg[i.rd] = reg[i.rt]>>reg[i.rs]; return 0;

				case F_ADD: reg[i.rd] = reg[i.rs] + reg[i.rt]; return 0;
				case F_ADDU: reg[i.rd] = (unsigned int)reg[i.rs] + (unsigned int)reg[i.rt]; return 0;
				case F_SUB: reg[i.rd] = reg[i.rs] - reg[i.rt]; return 0;
				case F_SUBU: reg[i.rd] = (unsigned int)reg[i.rs] - (unsigned int)reg[i.rt]; return 0;
				
				case F_AND: reg[i.rd] = reg[i.rs] & reg[i.rt]; return 0;
				case F_OR: reg[i.rd] = reg[i.rs] | reg[i.rt]; return 0;
				case F_XOR: reg[i.rd] = reg[i.rs] ^ reg[i.rt]; return 0;
				case F_NOR: reg[i.rd] = ~(reg[i.rs] | reg[i.rt]); return 0;

				case F_SLT: reg[i.rd] = reg[i.rs] < reg[i.rt] ? 1 : 0; return 0;
				case F_SLTU: reg[i.rd] = (unsigned int)reg[i.rs] < (unsigned int)reg[i.rt] ? 1 : 0; return 0;
			}
			return 1;
		}
	}

	return 1;
}

enum {
	MEMOP_NONE = 0,
	MEMOP_READ = 1,
	MEMOP_WRITE = 2,
};
unsigned int mem_addr = 0;
unsigned char mem_size = 0;
unsigned char mem_op = 0;

int mem(instr i) {
	unsigned int addr;
	switch(i.opcode) {
		// LB and LBU is implemented, but is not part of smips implementation...
		case LB:
			addr = reg[i.rs] + i.imm;
			reg[i.rt] = (signed char)mem_LB(addr);
			mem_op = MEMOP_READ;
			mem_addr = addr;
			mem_size = 1;
			return 0;	
		case LBU:
			addr = reg[i.rs] + i.imm;
			reg[i.rt] = (unsigned char)mem_LB(addr);
			mem_op = MEMOP_READ;
			mem_addr = addr;
			mem_size = 1;
			return 0;
		case LH:
			addr = reg[i.rs] + i.imm;
			reg[i.rt] = (signed short)mem_LHU(addr);
			mem_op = MEMOP_READ;
			mem_addr = addr;
			mem_size = 2;
			return 0;	
		case LHU:
			addr = reg[i.rs] + i.imm;
			reg[i.rt] = mem_LHU(addr);
			mem_op = MEMOP_READ;
			mem_addr = addr;
			mem_size = 2;
			return 0;	
		case LW: 
			addr = reg[i.rs] + i.imm;
			reg[i.rt] = mem_LW(addr);

			mem_op = MEMOP_READ;
			mem_addr = addr;
			mem_size = 4;

			return 0;
		case SB:
			addr = reg[i.rs] + i.imm;
			mem_SB(addr, reg[i.rt]);
			mem_op = MEMOP_WRITE;
			mem_addr = addr;
			mem_size = 1;
			return 0;
		case SH:
			addr = reg[i.rs] + i.imm;
			mem_SH(addr, reg[i.rt]);
			mem_op = MEMOP_WRITE;
			mem_addr = addr;
			mem_size = 2;
			return 0;
		case SW:
			addr = reg[i.rs] + i.imm;
			mem_SW(addr, reg[i.rt]);
			
			mem_op = MEMOP_WRITE;
			mem_addr = addr;
			mem_size = 4;

			return 0;
	}

	return 1;
}

///////////////////////////////////////////////
// Helper functions

int reg_old[32];
void backup_oldstate() {
	memcpy(reg_old, reg, sizeof(reg_old));
}
unsigned int atox(char* c) {
	unsigned int v = 0;
	while (*c) {
		if ( *c >= 'a' && *c <= 'z' ) {
			v *= 0x10;
			v += *c - 'a' + 0xa;
		} else if ( *c >= 'A' && *c <= 'Z' ) {
			v *= 0x10;
			v += *c - 'A' + 0xa;
		} else if ( *c >= '0' && *c <= '9' ) {
			v *= 0x10;
			v += *c - '0';
		} else {
			return v;
		}
		c += 1;
	}
	return v;
}

//////////////////////////////////////////////////////////////////////
// Single-stepping implementation

#define HASH_SIZE 1024 
//total meaningful bits: 1 + 2 + 2 + 5 = 10 => 1024 entries

typedef struct {
	unsigned int flag;
	unsigned int addr;
	unsigned int value;
	unsigned int pc;

	unsigned int addrd;
	unsigned int valued;
	unsigned int pcd;
} type_last;
type_last* type_hash = 0;

struct opcount {
	unsigned char opcode;
	unsigned char fn;
	unsigned int count;
};

struct opcount opcount_data[1024*4];
unsigned int opcount_count = 0;

unsigned long long total_steps = 0;
unsigned long long total_matches = 0;

unsigned int match_count = 0;

unsigned int pred_pc = 0;
unsigned int pred_branch_delay = 0;


unsigned char verbose = 0;
//FILE* fetchlog = 0;
//FILE* fdecompressed = 0;
void step_fullog(FILE* foutput) {
	char* l = log_get_nextline();
	while( l[0] != '#' ) {
		log_read_nextline();
		l = log_get_nextline();
	}
	//////////////////////////////////////////////////////////////////
	// Read single step log from file

	// l[0] is # at this point
	// instruction counter
	unsigned int ic = atox(l + 1);
	unsigned int pc = 0;
	int rreg[32] = {0,};
	int rvalid[32] = {0,};

	unsigned int memread_addr = 0;
	unsigned int memread_value = 0;
	unsigned char memread_size = 0;

	unsigned int memwrite_addr = 0;
	unsigned int memwrite_value = 0;
	unsigned char memwrite_size = 0;

	// l contains first line after # at this point
	log_read_nextline();
	l = log_get_nextline();

	while( l[0] != '#' ) {

		if (l[0] == 'S') { // STEP
			pc = atox(l + 4);
		} else if (l[0] == 'R') { // REG
			int regnum = atoi(l + 4);
			int shnum = atoi(l + 7); // shadow registers not used...
			unsigned int val = atox(l + 10);
			rreg[regnum] = val;
			rvalid[regnum] = 1;
		} else if (l[4] == 'R') { // MEM_READ
			memread_addr = atox(l + 9);
			memread_size = atoi(l + 22);
			memread_value = atox(l + 25);
		} else if (l[4] == 'W') { // MEM_WRITE
			memwrite_addr = atox(l + 10);
			memwrite_size = atoi(l + 23);
			memwrite_value = atox(l + 26);
		} else {
			fprintf(stderr, "error: unrecognized log entry!\n" );
		}

		log_read_nextline();
		l = log_get_nextline();
	}
	
	//////////////////////////////////////////////////////////////////
	// Single-step predictor processor

	mem_op = MEMOP_NONE;
	pred_pc = (pred_pc >> 2) << 2;
	backup_oldstate();
	mem_clear();
	instr inst;
	
	// fetch
	unsigned int fetched_pc = pred_pc;
	unsigned int pred_ic = fetch(pred_pc);

	//decode
	decode(pred_ic, &inst);

	//execute
	reg[0] = 0;
	int ret = mem(inst);
	if ( ret ) {
		ret = alu(inst);
	}
	
	if ( pred_branch_delay ) {
		pred_pc = pred_branch_delay;
		pred_branch_delay = 0;
	} else {
		if ( ret ) pred_branch_delay = next_pc(pred_pc, inst);
		pred_pc = pred_pc + 4;
	}

	if ( ret && !pred_branch_delay && !next_pc(pred_pc, inst) ) {
		int i;
		for ( i = 0; i < opcount_count; i++ ) {
			if ( opcount_data[i].opcode == inst.opcode /*&& opcount_data[i].fn == inst.fn*/ ) {
				opcount_data[i].count ++;
				break;
			}
		}
		if ( i >= opcount_count && opcount_count < 1024*4) {
			opcount_data[opcount_count].opcode = inst.opcode;
			opcount_data[opcount_count].fn = inst.fn;
			opcount_data[opcount_count++].count = 1;
		}
	}


	//////////////////////////////////////////////////////////////////
	// Compare results

	int different = 0;
	if ( pred_pc != pc ) {
		different = 1; 
	}
/*
	if (pred_ic != ic) {
		different = 2;
	}
*/
	int i;
	for ( i = 0; i < 32; i++ ) {
		if (rvalid[i]) {
			if ( rreg[i] != reg[i] ) {
				different = 100 + i;
			}
		} else {
			if ( reg[i] != reg_old[i] ) {
				different = 200 + i;
			}
		}
	}

	if ( mem_op == MEMOP_READ && !memread_size ) {
		different = 13;
	}
	if ( mem_op == MEMOP_WRITE && !memwrite_size ) different = 14;

	if (memread_size) {
		if (mem_op !=MEMOP_READ) different = 3;
		else {
			if (mem_addr != memread_addr) {
				different = 3;
			}
			if (mem_size != memread_size) {
				different = 3;
			}
		}
	}
	if (memwrite_size) {
		if (mem_op !=MEMOP_WRITE) different = 4;
		else {
			if (mem_addr != memwrite_addr) {
				different = 4;
			}
			if (mem_size != memwrite_size) {
				different = 4;
			}
		}
	}
	
	//////////////////////////////////////////////////////////////////////
	// Update mismatches

	total_steps += 1;
	
	if ( !different ) {
		match_count += 1;
		total_matches += 1;
		mem_commit();
		
//		if ( !fetchlog ) fetchlog = fopen("fetchlog.txt", "w");
//		fprintf(fetchlog, "Fetched %08x(%08x)\n", fetched_pc, pred_ic );

	} else {
		// dump match count;

		if (match_count >= (1<<16)) {
			//printf( "matches! %d\n", match_count );
			unsigned int t = match_count % (1<<30);
			t |= 1<<30;
//			if ( match_count >= (1<<16) ) {
				fwrite(&t, 4, 1, foutput);
				match_count = 0;
//			}

			if ( verbose )
				printf( "%15s:%d\n", "match", match_count );
		}

		//match_count = 0;
		//printf ( "%x %x mismatch[%x]! %d\n", pc, pred_pc, ic, different );
		fflush(stdout);


		unsigned int flags = 0;
		unsigned int addr = 0;
		unsigned int value = 0;
		
		unsigned char target = 0;
		unsigned char size = 0;
		unsigned char memAcc = 0;
		unsigned char regChange = 0;

		pred_pc = pc;
//		if ( !fetchlog ) fetchlog = fopen("fetchlog.txt", "w");
//		fprintf(fetchlog, "Fetched feed %08x\n", pc );

		pred_branch_delay = 0; // is this necessary?
		// take care of mem_SW and LW... later
		if (memread_size) {
			mem_SW(memread_addr, memread_value);
			size = memread_size;
			memAcc = 1;
			addr = memread_addr;
			value = memread_value;
			
		}
		if (memwrite_size) {
			mem_SW(memwrite_addr, memwrite_value);
			size = memwrite_size;
			memAcc = 2;
			addr = memwrite_addr;
			value = memwrite_value;

		}
		int i;
		for ( i = 0; i < 32; i++ ) {
			if (rvalid[i]) {
				reg[i] = rreg[i];
	
				target = i;
				value = rreg[i];
				regChange = 1;

			} else if (reg[i] != reg_old[i]) {
				reg[i] = reg_old[i];
			}
		}

		if ( size > 4 ) {
			printf( "size > 4?!?!\n" );
		}
		
		if ( size == 4 ) size = 3;
		unsigned int idx = (regChange%2) + ((memAcc%4)<<1) + ((size%4)<<3) + ((target%32)<<5);
		type_last* hashv = &type_hash[idx%HASH_SIZE]; // mod by HASH_SIZE just in case
		int addrd = (int)addr - hashv->addr;
		int valued = (int)value - hashv->value;
		int pcd = (int)pc - hashv->pc;
		int addrdd = addrd-hashv->addrd;
		int valuedd = valued - hashv->valued;
		int pcdd = pcd - hashv->pcd;
		
		flags = (regChange<<16) + (memAcc<<17)+((size%(1<<2))<<19)
			+(target<<24);
		//printf( "%x--\n", flags );
		// flags[21:23] are reserved for second order compressions
		if ( match_count ) flags += match_count;
		match_count = 0;
		
		if ( memAcc && !addrdd ) {
			flags |= (1<<21);
		}
		if ( regChange && !valuedd ) {
			flags |= (1<<22);
		}
		if ( !pcdd ) {
			flags |= (1<<23);
		}


		fwrite(&flags, 4, 1, foutput);
		if ( memAcc && addrdd ) fwrite(&addr, 4, 1, foutput);
		if ( /*memAcc ||*/ regChange && valuedd ) fwrite(&value, 4, 1, foutput);
		if ( pcdd )
			fwrite(&pc, 4, 1, foutput);
	
		if ( verbose ) {
			if ( memAcc ) printf( "%15s:%x %x %x %x\n", "memacc", flags, addr, value, pc );
			else if ( regChange ) printf( "%15s:%x %x %x\n", "regchange", flags, value, pc );
			else printf( "%15s:%x %x\n", "other", flags, pc );
		}

		mem_commit();
		
		hashv->addr = addr;
		hashv->value = value;
		hashv->pc = pc;

		hashv->addrd = addrd;
		hashv->valued = valued;
		hashv->pcd = pcd;
	}
}

void traverse_fullog(unsigned int entrypoint, unsigned int steps) {
	pred_pc = entrypoint;
	FILE* fo = fopen("compressed.bin", "wb");

	log_read_nextline();
	char* l = log_get_nextline();
	unsigned long step_count = 0;
	while ( l[1] && (!steps || step_count < steps )) { // log_get_nextline returns "#\0" at eof
		step_fullog(fo);

		step_count += 1;
		l = log_get_nextline();
	}
	fclose(fo);
	printf( "Done traversing!\n");
	printf( "%llu/%llu\n", total_matches, total_steps );
}

extern unsigned int total_loads;
extern unsigned int load_misses;
extern unsigned int total_writes;
extern unsigned int write_misses;
int main(int argc, char** argv) {
	if ( argc < 2 ) {
		fprintf(stderr, "usage: %s [filename]\n", argv[0] );
		exit(-1);
	}
	total_steps = 0;
	total_matches = 0;
	opcount_count = 0;

	//TODO: sanity check on argvs?
	int step = 0;
	if ( argc >= 3 ) {
		step = atoi(argv[2]);
	}
	if ( argc >= 4 ) {
		char* flags = 0;
		flags = argv[3];
		int i;
		for ( i = 0; i < strlen(flags); i++ ) {
			if ( flags[i] == 'v' ) verbose = 1;
		}
	}

	FILE* bin = fopen(argv[1], "rb");

	if ( !bin ) {
		fprintf(stderr, "error: file %s not found\n", argv[1] );
		exit(-1);
	}

	int i;
	elf_load_header(bin);
	unsigned int entrypoint = elf_entrypoint(bin);
	printf( "entry: %x\n", entrypoint );

	esections = elf_load_exec_section(bin);
	while ( esections[++ecount].data );
	
	printf( "executable segment count: %d\n", ecount );

	mem_init();
	full_log_init(argv[1]);
	
	type_hash = (type_last*)calloc(HASH_SIZE, sizeof(type_last));
	memset(type_hash, 0, sizeof(type_last)*HASH_SIZE);

	if ( !logfile ) {
		fprintf(stderr, "error: %s.log not found!!\n", argv[1] );
		exit(-1);
	}

	printf( "-------\n\n" );
	FILE* foutput = fopen("exec.bin", "wb");
	fwrite(&entrypoint, sizeof(unsigned int), 1, foutput);
	for (i = 0; i < ecount; i++) {
		fwrite(&esections[i].offset, sizeof(unsigned int), 1, foutput);
		fwrite(&esections[i].size, sizeof(unsigned int), 1, foutput);
		fwrite(esections[i].data, esections[i].size, 1, foutput);
	}
	fclose(foutput);

	traverse_fullog(entrypoint, step);
	printf( "L:%d miss/%d, W:%d miss/%d\n\n", load_misses, total_loads, write_misses, total_writes );

	unsigned long long nonimpl = 0;
	for ( i = 0; i < opcount_count; i++ ) {
		nonimpl += opcount_data[i].count;
	}
	for ( i = 0; i < opcount_count; i++ ) {
		int j;
		for ( j = i; j < opcount_count; j ++ ) {
			if ( opcount_data[i].count > opcount_data[j].count ) {
				unsigned int c = opcount_data[i].count;
				opcount_data[i].count = opcount_data[j].count;
				opcount_data[j].count = c;
				unsigned char cc = opcount_data[i].opcode;
				opcount_data[i].opcode = opcount_data[j].opcode;
				opcount_data[j].opcode = cc;
				cc = opcount_data[i].fn;
				opcount_data[i].fn = opcount_data[j].fn;
				opcount_data[j].fn = cc;
			}
		}
	}
	printf( "Nonimplemented: %llu\n", nonimpl );
	for ( i = 0; i < opcount_count; i++ ) {
		printf( "%x %x %d\n", opcount_data[i].opcode, opcount_data[i].fn, opcount_data[i].count );
	}

	exit(0);
}
